package com.whut.imservice.ws.handler;

import cn.hutool.core.lang.Pair;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.whut.imservice.constant.MessageTypeEnum;
import com.whut.imservice.domain.vo.MessageVO;
import com.whut.imservice.factory.EnhancedThreadFactory;
import com.whut.imservice.service.MsgService;
import io.netty.channel.*;
import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
import io.netty.handler.codec.http.websocketx.WebSocketFrame;
import io.netty.util.AttributeKey;
import jakarta.annotation.PreDestroy;
import lombok.RequiredArgsConstructor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Optional;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicLong;

/**
 * @Description: 服务端处理消息的handler
 * @Author: Zonghe Wang
 * @Date: 2025/5/24 20:06
 */
@ChannelHandler.Sharable
@Component
@RequiredArgsConstructor
public class WebsocketRouterHandler extends SimpleChannelInboundHandler<WebSocketFrame> {

    private static final Logger log = LoggerFactory.getLogger(WebsocketRouterHandler.class);
    private static final int MAX_RETRY_TIMES = 2;
    private static final int ACK_CHECK_DELAY = 5;
    private static final int RETRY_INTERVAL = 2;

    // channel attributes
    private static final AttributeKey<AtomicLong> TID_GENERATOR = AttributeKey.valueOf("tidGenerator");
    private static final AttributeKey<ConcurrentHashMap<Long, JSONObject>> NON_ACKED_MAP = AttributeKey.valueOf("nonAckedMap");
    private static final AttributeKey<Long> USER_ID_ATTR = AttributeKey.valueOf("userId");

    // 连接channel管理
    private static final ConcurrentHashMap<Long, Channel> userChannels = new ConcurrentHashMap<>(15000);
    private static final ConcurrentHashMap<Channel, Long> channelUsers = new ConcurrentHashMap<>(15000);

    // ACK
    private final ScheduledExecutorService ackExecutor = new ScheduledThreadPoolExecutor(
            50,
            new EnhancedThreadFactory("ACKChecker"),
            new ThreadPoolExecutor.CallerRunsPolicy()
    );

    private final MsgService msgService;

    @PreDestroy
    public void shutdown() {
        ackExecutor.shutdown();
        log.info("ACKChecker shutdown");
    }

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, WebSocketFrame webSocketFrame) {
        if (!(webSocketFrame instanceof TextWebSocketFrame)) {
            log.info("不支持除文本外类型:{}", webSocketFrame.getClass().getName());
            return;
        }

        try {
            String text = ((TextWebSocketFrame) webSocketFrame).text();
            JSONObject message = JSONObject.parseObject(text);
            MessageTypeEnum type = MessageTypeEnum.fromCode(message.getIntValue("type"));
            JSONObject data = message.getJSONObject("data");

            switch (type) {
                case HEARTBEAT -> handleHeartbeat(ctx, data);
                case LOGIN -> handleLogin(ctx, data);
                case QUERY_MESSAGES -> handleQueryMessages(ctx, data);
                case SEND_MESSAGE -> handleSendMessage(ctx, data);
                case QUERY_UNREAD -> handleQueryUnread(ctx, data);
                case ACK -> handleAck(ctx, data);
                default -> log.warn("不支持的消息类型:{}", type);
            }
        } catch (Exception e) {
            log.error("解析消息失败:{}", e.getMessage());
            sendErrorResponse(ctx);
        }
    }

    private void handleAck(ChannelHandlerContext ctx, JSONObject data) {
        Long tid = data.getLong("tid");
        getNonAckedMap(ctx.channel()).remove(tid);
        log.debug("接收到ACK - TID: {}", tid);
    }

    private void handleQueryUnread(ChannelHandlerContext ctx, JSONObject data) {
        Long uid = data.getLong("uid");
        Long unreadCount = msgService.queryTotalUnread(uid);
        sendResponse(ctx, MessageTypeEnum.QUERY_UNREAD, Pair.of("unread", unreadCount));
    }

    private void handleSendMessage(ChannelHandlerContext ctx, JSONObject data) {
        MessageVO message = msgService.sendNewMsg(
                data.getLong("senderUid"),
                data.getLong("recipientUid"),
                data.getString("content"),
                data.getIntValue("msgType")
        );
        if (message != null) {
            sendResponse(ctx, MessageTypeEnum.SEND_MESSAGE, Pair.of("message", JSON.toJSONString(message)));
        }
    }

    private void handleQueryMessages(ChannelHandlerContext ctx, JSONObject data) {
        Long ownerUid = data.getLong("ownerUid");
        Long otherUid = data.getLong("otherUid");
        List<MessageVO> messages = msgService.queryConversationMsg(ownerUid, otherUid);
        sendResponse(ctx, MessageTypeEnum.QUERY_MESSAGES, Pair.of("messages", messages));
    }

    private void handleLogin(ChannelHandlerContext ctx, JSONObject data) {
        Long uid = data.getLong("uid");
        Channel channel = ctx.channel();

        // 处理重复登录
        Channel existingChannel = userChannels.put(uid, channel);
        Optional.ofNullable(existingChannel)
                .filter(Channel::isActive)
                .ifPresent(c -> {
                    log.info("用户{}已登录，将断开旧连接", uid);
                    c.close();
                });

        channelUsers.put(channel, uid);
        channel.attr(USER_ID_ATTR).set(uid);
        initChannelAttributes(channel);

        log.info("用户{}登录成功, Channel: {}", uid, channel);
        sendResponse(ctx, MessageTypeEnum.LOGIN, Pair.of("status", "success"));
    }

    private void handleHeartbeat(ChannelHandlerContext ctx, JSONObject data) {
        Long uid = data.getLong("uid");
        Long timeout = data.getLong("timeout");
        log.debug("收到心跳包：uid:{}, timeout:{}", uid, timeout);
        sendResponse(ctx, MessageTypeEnum.HEARTBEAT, Pair.of("timeout", timeout));
    }

    // 工具方法 ----------------------------------------
    private ConcurrentHashMap<Long, JSONObject> getNonAckedMap(Channel channel) {
        return channel.attr(NON_ACKED_MAP).get();
    }
    private void sendResponse(ChannelHandlerContext ctx, MessageTypeEnum type, Pair<String, Object> data) {
        JSONObject response = new JSONObject();
        response.put("type", type.getCode());
        response.put("data", new JSONObject().fluentPut(data.getKey(), data.getValue()));
        ctx.writeAndFlush(new TextWebSocketFrame(JSON.toJSONString(response)));
    }

    private void sendErrorResponse(ChannelHandlerContext ctx) {
        JSONObject response = new JSONObject();
        response.put("type", MessageTypeEnum.ERROR.getCode());
        response.put("data", new JSONObject().fluentPut("error", "解析消息失败"));
        ctx.writeAndFlush(new TextWebSocketFrame(JSON.toJSONString(response)));
    }

    private void initChannelAttributes(Channel channel) {
        channel.attr(TID_GENERATOR).set(new AtomicLong(0));
        channel.attr(NON_ACKED_MAP).set(new ConcurrentHashMap<>());
    }

    // 连接管理 ----------------------------------------
    @Override
    public void channelInactive(ChannelHandlerContext ctx) {
        Channel channel = ctx.channel();
        Long uid = channelUsers.remove(channel);
        if (uid != null) {
            userChannels.remove(uid);
            log.info("Channel 关闭 - UID: {}, Channel: {}", uid, channel);
        }
        cleanupChannel(channel);
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
        Channel channel = ctx.channel();
        Long uid = channelUsers.get(channel);
        log.error("Channel 异常 - UID: {}, Channel: {}", uid, channel, cause);
        cleanupChannel(channel);
        ctx.close();
    }

    public void cleanupChannel(Channel channel) {
        channel.attr(NON_ACKED_MAP).get().clear();
        channel.attr(TID_GENERATOR).set(null);
    }

    // 消息推送 ----------------------------------------
    public void pushMessage(long recipientUid, JSONObject message) {
        Optional.ofNullable(userChannels.get(recipientUid))
                .filter(channel -> channel.isActive() && channel.isWritable())
                .ifPresentOrElse(
                        channel -> sendMessageWithRetry(channel, message),
                        () -> log.warn("User offline or channel unavailable - UID: {}", recipientUid)
                );
    }

    private void sendMessageWithRetry(Channel channel, JSONObject message) {
        long tid = generateTid(channel);
        message.put("tid", tid);

        ChannelFuture future = channel.writeAndFlush(new TextWebSocketFrame(message.toJSONString()));
        future.addListener((ChannelFutureListener) f -> {
            if (f.isSuccess()) {
                scheduleAckCheck(channel, message, tid);
            } else {
                log.error("Message send failed - TID: {}", tid, f.cause());
            }
        });
    }

    private long generateTid(Channel channel) {
        return channel.attr(TID_GENERATOR).get().incrementAndGet();
    }

    private void scheduleAckCheck(Channel channel, JSONObject message, long tid) {
        getNonAckedMap(channel).put(tid, message);

        ackExecutor.schedule(() -> {
            if (!channel.isActive()) return;
            if (getNonAckedMap(channel).containsKey(tid)) {
                retrySendMessage(channel, message, tid, MAX_RETRY_TIMES);
            }
        }, ACK_CHECK_DELAY, TimeUnit.SECONDS);
    }

    private void retrySendMessage(Channel channel, JSONObject message, long tid, int retries) {
        if (retries <= 0) {
            log.warn("Max retries exceeded - TID: {}", tid);
            getNonAckedMap(channel).remove(tid);
            return;
        }

        channel.writeAndFlush(new TextWebSocketFrame(message.toJSONString()))
                .addListener((ChannelFutureListener) future -> {
                    if (future.isSuccess()) {
                        ackExecutor.schedule(
                                () -> checkAndRetry(channel, tid, retries),
                                RETRY_INTERVAL,
                                TimeUnit.SECONDS
                        );
                    }
                });
    }

    private void checkAndRetry(Channel channel, long tid, int retries) {
        if (channel.isActive() && getNonAckedMap(channel).containsKey(tid)) {
            log.warn("Retrying message - TID: {}, Retries left: {}", tid, retries - 1);
            retrySendMessage(channel, getNonAckedMap(channel).get(tid), tid, retries - 1);
        }
    }

}
